package com.leon.datalink.core.persistence.plugin;

import akka.persistence.snapshot.japi.SnapshotStore;
import akka.persistence.SelectedSnapshot;
import akka.persistence.SnapshotMetadata;
import akka.persistence.SnapshotSelectionCriteria;
import com.leon.datalink.core.evn.EnvUtil;
import com.alibaba.druid.pool.DruidDataSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.concurrent.Future;
import akka.dispatch.Futures;

import java.io.*;
import java.sql.*;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.Map;

/**
 * 基于PostgreSQL的Snapshot插件
 * 使用PostgreSQL数据库进行真实的快照持久化
 * 
 * @author DataLink Team
 * @version 1.0.0
 */
public class PostgresSnapshotPlugin extends SnapshotStore {
    
    private static final Logger logger = LoggerFactory.getLogger(PostgresSnapshotPlugin.class);
    
    // PostgreSQL数据源
    private DruidDataSource dataSource;
    
    // 表名
    private static final String SNAPSHOT_TABLE = "datalink_snapshots";
    
    // 内存索引（用于快速查找）
    private final Map<String, SnapshotMetadata> metadataIndex = new ConcurrentHashMap<>();
    
    public PostgresSnapshotPlugin() {
        // Akka会自动调用这个构造函数
    }
    
    @Override
    public void preStart() throws Exception {
        super.preStart();
        
        // 初始化PostgreSQL数据源
        initializeDataSource();
        
        // 创建表结构
        createTablesIfNotExists();

        // 加载现有的快照元数据索引
        loadMetadataIndex();

        logger.info("PostgreSQL Snapshot Plugin initialized successfully");
    }
    
    /**
     * 初始化PostgreSQL数据源
     */
    private void initializeDataSource() throws Exception {
        String host = System.getProperty("datalink.persistence.postgres.host", "localhost");
        String port = System.getProperty("datalink.persistence.postgres.port", "5432");
        String database = System.getProperty("datalink.persistence.postgres.database", "datalink");
        String username = System.getProperty("datalink.persistence.postgres.username", "datalink");
        String password = System.getProperty("datalink.persistence.postgres.password", "datalink");
        
        String url = String.format("jdbc:postgresql://%s:%s/%s?useUnicode=true&characterEncoding=utf-8&useSSL=false",
                host, port, database);
        
        dataSource = new DruidDataSource();
        dataSource.setDriverClassName("org.postgresql.Driver");
        dataSource.setUrl(url);
        dataSource.setUsername(username);
        dataSource.setPassword(password);
        dataSource.setInitialSize(5);
        dataSource.setMinIdle(2);
        dataSource.setMaxActive(20);
        dataSource.setValidationQuery("SELECT 1");
        dataSource.setTestWhileIdle(true);
        dataSource.setTestOnBorrow(false);
        dataSource.setTestOnReturn(false);
        
        // 初始化连接池
        dataSource.init();
        
        logger.info("PostgreSQL DataSource initialized: {}", url);
    }
    
    /**
     * 创建表结构
     */
    private void createTablesIfNotExists() throws Exception {
        String createTableSQL = String.format(
            "CREATE TABLE IF NOT EXISTS %s (" +
            "id BIGSERIAL PRIMARY KEY," +
            "persistence_id VARCHAR(255) NOT NULL," +
            "sequence_nr BIGINT NOT NULL," +
            "timestamp_ms BIGINT NOT NULL," +
            "snapshot_data BYTEA NOT NULL," +
            "created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP," +
            "UNIQUE(persistence_id, sequence_nr)" +
            ")", SNAPSHOT_TABLE);

        String createIndexSQL = String.format(
            "CREATE INDEX IF NOT EXISTS idx_%s_persistence_id_sequence " +
            "ON %s (persistence_id, sequence_nr DESC)", SNAPSHOT_TABLE, SNAPSHOT_TABLE);

        String createTimestampIndexSQL = String.format(
            "CREATE INDEX IF NOT EXISTS idx_%s_persistence_id_timestamp " +
            "ON %s (persistence_id, timestamp_ms DESC)", SNAPSHOT_TABLE, SNAPSHOT_TABLE);
        
        try (Connection conn = dataSource.getConnection();
             Statement stmt = conn.createStatement()) {
            
            stmt.execute(createTableSQL);
            stmt.execute(createIndexSQL);
            stmt.execute(createTimestampIndexSQL);
            
            logger.info("Snapshot table {} created/verified", SNAPSHOT_TABLE);
        }
    }
    
    @Override
    public Future<Optional<SelectedSnapshot>> doLoadAsync(String persistenceId,
                                                        SnapshotSelectionCriteria criteria) {
        try {
            // 查找符合条件的最新快照
            SelectedSnapshot selectedSnapshot = findLatestSnapshotFromDatabase(persistenceId, criteria);

            if (selectedSnapshot != null) {
                logger.debug("Loaded snapshot for persistenceId: {} at sequence: {}",
                           persistenceId, selectedSnapshot.metadata().sequenceNr());
                return Futures.successful(Optional.of(selectedSnapshot));
            } else {
                logger.debug("No snapshot found for persistenceId: {}", persistenceId);
                return Futures.successful(Optional.empty());
            }
        } catch (Exception e) {
            logger.error("Failed to load snapshot for persistenceId: {}", persistenceId, e);
            return Futures.successful(Optional.empty());
        }
    }
    
    @Override
    public Future<Void> doSaveAsync(SnapshotMetadata metadata, Object snapshot) {
        try {
            // 保存快照到数据库
            saveSnapshotToDatabase(metadata, snapshot);

            // 更新内存索引
            String key = createSnapshotKey(metadata.persistenceId(), metadata.sequenceNr());
            metadataIndex.put(key, metadata);

            logger.debug("Saved snapshot for persistenceId: {} at sequence: {}",
                       metadata.persistenceId(), metadata.sequenceNr());
            return Futures.successful(null);
        } catch (Exception e) {
            logger.error("Failed to save snapshot for persistenceId: {}",
                       metadata.persistenceId(), e);
            return Futures.failed(e);
        }
    }
    
    @Override
    public Future<Void> doDeleteAsync(SnapshotMetadata metadata) {
        try {
            // 删除快照记录
            String deleteSQL = "DELETE FROM " + SNAPSHOT_TABLE + 
                             " WHERE persistence_id = ? AND sequence_nr = ?";
            
            try (Connection conn = dataSource.getConnection();
                 PreparedStatement stmt = conn.prepareStatement(deleteSQL)) {
                
                stmt.setString(1, metadata.persistenceId());
                stmt.setLong(2, metadata.sequenceNr());
                
                int deletedCount = stmt.executeUpdate();
                
                // 从索引中移除
                String key = createSnapshotKey(metadata.persistenceId(), metadata.sequenceNr());
                metadataIndex.remove(key);

                logger.debug("Deleted {} snapshot for persistenceId: {} at sequence: {}",
                           deletedCount, metadata.persistenceId(), metadata.sequenceNr());
            }
            
            return Futures.successful(null);
        } catch (Exception e) {
            logger.error("Failed to delete snapshot for persistenceId: {}",
                       metadata.persistenceId(), e);
            return Futures.failed(e);
        }
    }
    
    @Override
    public Future<Void> doDeleteAsync(String persistenceId, SnapshotSelectionCriteria criteria) {
        try {
            // 删除符合条件的快照记录
            deleteSnapshotsByCriteriaFromDatabase(persistenceId, criteria);

            logger.debug("Deleted snapshots for persistenceId: {} with criteria: {}",
                       persistenceId, criteria);
            return Futures.successful(null);
        } catch (Exception e) {
            logger.error("Failed to delete snapshots for persistenceId: {}", persistenceId, e);
            return Futures.failed(e);
        }
    }
    
    /**
     * 将快照保存到数据库
     */
    private void saveSnapshotToDatabase(SnapshotMetadata metadata, Object snapshot) throws Exception {
        String insertSQL = "INSERT INTO " + SNAPSHOT_TABLE + 
                         " (persistence_id, sequence_nr, timestamp_ms, snapshot_data) " +
                         "VALUES (?, ?, ?, ?) ON CONFLICT (persistence_id, sequence_nr) " +
                         "DO UPDATE SET timestamp_ms = EXCLUDED.timestamp_ms, snapshot_data = EXCLUDED.snapshot_data";
        
        try (Connection conn = dataSource.getConnection();
             PreparedStatement stmt = conn.prepareStatement(insertSQL)) {
            
            stmt.setString(1, metadata.persistenceId());
            stmt.setLong(2, metadata.sequenceNr());
            stmt.setLong(3, metadata.timestamp());
            
            // 序列化快照数据
            byte[] snapshotBytes = serializeSnapshot(snapshot);
            stmt.setBytes(4, snapshotBytes);
            
            stmt.executeUpdate();
            
            logger.debug("Saved snapshot to database for persistenceId: {} at sequence: {}",
                       metadata.persistenceId(), metadata.sequenceNr());
        }
    }
    
    /**
     * 从数据库查找最新快照
     */
    private SelectedSnapshot findLatestSnapshotFromDatabase(String persistenceId, SnapshotSelectionCriteria criteria) {
        try {
            String selectSQL = "SELECT persistence_id, sequence_nr, timestamp_ms, snapshot_data " +
                             "FROM " + SNAPSHOT_TABLE + 
                             " WHERE persistence_id = ? AND sequence_nr >= ? AND sequence_nr <= ? " +
                             "AND timestamp_ms >= ? AND timestamp_ms <= ? " +
                             "ORDER BY sequence_nr DESC LIMIT 1";
            
            try (Connection conn = dataSource.getConnection();
                 PreparedStatement stmt = conn.prepareStatement(selectSQL)) {
                
                stmt.setString(1, persistenceId);
                stmt.setLong(2, criteria.minSequenceNr());
                stmt.setLong(3, criteria.maxSequenceNr());
                stmt.setLong(4, criteria.minTimestamp());
                stmt.setLong(5, criteria.maxTimestamp());
                
                try (ResultSet rs = stmt.executeQuery()) {
                    if (rs.next()) {
                        String readPersistenceId = rs.getString("persistence_id");
                        long sequenceNr = rs.getLong("sequence_nr");
                        long timestamp = rs.getLong("timestamp_ms");
                        byte[] snapshotBytes = rs.getBytes("snapshot_data");
                        
                        // 反序列化快照数据
                        Object snapshot = deserializeSnapshot(snapshotBytes);
                        
                        SnapshotMetadata metadata = SnapshotMetadata.apply(readPersistenceId, sequenceNr, timestamp);
                        return new SelectedSnapshot(metadata, snapshot);
                    }
                }
            }
        } catch (Exception e) {
            logger.error("Failed to find latest snapshot from database", e);
        }
        
        return null;
    }
    
    /**
     * 根据条件删除快照记录
     */
    private void deleteSnapshotsByCriteriaFromDatabase(String persistenceId, SnapshotSelectionCriteria criteria) {
        try {
            String deleteSQL = "DELETE FROM " + SNAPSHOT_TABLE + 
                             " WHERE persistence_id = ? AND sequence_nr >= ? AND sequence_nr <= ? " +
                             "AND timestamp_ms >= ? AND timestamp_ms <= ?";
            
            try (Connection conn = dataSource.getConnection();
                 PreparedStatement stmt = conn.prepareStatement(deleteSQL)) {
                
                stmt.setString(1, persistenceId);
                stmt.setLong(2, criteria.minSequenceNr());
                stmt.setLong(3, criteria.maxSequenceNr());
                stmt.setLong(4, criteria.minTimestamp());
                stmt.setLong(5, criteria.maxTimestamp());
                
                int deletedCount = stmt.executeUpdate();
                
                // 从内存索引中移除相关条目
                metadataIndex.entrySet().removeIf(entry -> {
                    String key = entry.getKey();
                    SnapshotMetadata metadata = entry.getValue();
                    return key.startsWith(persistenceId + "-snapshot-") &&
                           metadata.sequenceNr() >= criteria.minSequenceNr() &&
                           metadata.sequenceNr() <= criteria.maxSequenceNr() &&
                           metadata.timestamp() >= criteria.minTimestamp() &&
                           metadata.timestamp() <= criteria.maxTimestamp();
                });
                
                logger.debug("Deleted {} snapshots from database for persistenceId: {}", 
                           deletedCount, persistenceId);
            }
        } catch (Exception e) {
            logger.error("Failed to delete snapshots by criteria from database", e);
        }
    }
    
    /**
     * 序列化快照
     */
    private byte[] serializeSnapshot(Object snapshot) throws Exception {
        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
             ObjectOutputStream oos = new ObjectOutputStream(baos)) {
            
            if (snapshot instanceof Serializable) {
                oos.writeObject(snapshot);
            } else {
                oos.writeObject(snapshot.toString());
            }
            
            oos.flush();
            return baos.toByteArray();
        }
    }
    
    /**
     * 反序列化快照
     */
    private Object deserializeSnapshot(byte[] snapshotBytes) throws Exception {
        try (ByteArrayInputStream bais = new ByteArrayInputStream(snapshotBytes);
             ObjectInputStream ois = new ObjectInputStream(bais)) {
            
            return ois.readObject();
        }
    }
    
    /**
     * 创建快照键
     */
    private String createSnapshotKey(String persistenceId, long sequenceNr) {
        return persistenceId + "-snapshot-" + String.format("%020d", sequenceNr);
    }
    
    /**
     * 加载元数据索引
     */
    private void loadMetadataIndex() {
        try {
            String selectSQL = "SELECT persistence_id, sequence_nr, timestamp_ms FROM " + SNAPSHOT_TABLE;
            
            try (Connection conn = dataSource.getConnection();
                 Statement stmt = conn.createStatement();
                 ResultSet rs = stmt.executeQuery(selectSQL)) {
                
                while (rs.next()) {
                    String persistenceId = rs.getString("persistence_id");
                    long sequenceNr = rs.getLong("sequence_nr");
                    long timestamp = rs.getLong("timestamp_ms");
                    
                    SnapshotMetadata metadata = SnapshotMetadata.apply(persistenceId, sequenceNr, timestamp);
                    String key = createSnapshotKey(persistenceId, sequenceNr);
                    metadataIndex.put(key, metadata);
                }
                
                logger.info("Loaded metadata index with {} entries", metadataIndex.size());
            }
        } catch (Exception e) {
            logger.error("Failed to load metadata index", e);
        }
    }
    
    @Override
    public void postStop() throws Exception {
        // 关闭数据源
        if (dataSource != null) {
            dataSource.close();
        }
        
        // 清理内存索引
        if (metadataIndex != null) {
            metadataIndex.clear();
        }
        
        logger.info("PostgreSQL Snapshot Plugin stopped");
        super.postStop();
    }
}
